#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/time.h>
#include "md5.h"
#include <stdint.h>

#define DEFAULT_IP "127.0.0.1"
#define DEFAULT_PORT 9001
#define CHUNK_SIZE 4096

void send_udp(const char *dest_ip, int dest_port, int local_port, unsigned char *buffer, size_t filesize, size_t *packet_count) {
    int sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock < 0) {
        perror("socket error");
        exit(1);
    }
    // 调大UDP缓冲区
    int bufsize = 8 * 1024 * 1024;
    setsockopt(sock, SOL_SOCKET, SO_SNDBUF, &bufsize, sizeof(bufsize));
    setsockopt(sock, SOL_SOCKET, SO_RCVBUF, &bufsize, sizeof(bufsize));
    if (local_port > 0) {
        struct sockaddr_in local_addr;
        memset(&local_addr, 0, sizeof(local_addr));
        local_addr.sin_family = AF_INET;
        local_addr.sin_addr.s_addr = htonl(INADDR_ANY);
        local_addr.sin_port = htons(local_port);
        if (bind(sock, (struct sockaddr*)&local_addr, sizeof(local_addr)) < 0) {
            perror("bind local port error");
            close(sock);
            exit(1);
        }
    }
    struct sockaddr_in dest_addr;
    memset(&dest_addr, 0, sizeof(dest_addr));
    dest_addr.sin_family = AF_INET;
    dest_addr.sin_port = htons(dest_port);
    inet_pton(AF_INET, dest_ip, &dest_addr.sin_addr);
    size_t total_sent = 0;
    unsigned int seq = 0;
    while (total_sent < filesize) {
        size_t chunk = (filesize - total_sent > CHUNK_SIZE) ? CHUNK_SIZE : (filesize - total_sent);
        unsigned char sendbuf[CHUNK_SIZE + 4];
        memcpy(sendbuf, &seq, 4);
        memcpy(sendbuf + 4, buffer + total_sent, chunk);
        ssize_t sent = sendto(sock, sendbuf, chunk + 4, 0, (struct sockaddr*)&dest_addr, sizeof(dest_addr));
        if (sent < 0) {
            perror("sendto error");
            break;
        }
        total_sent += chunk;
        seq++;
        (*packet_count)++;
        usleep(1); // 适当限速，降低丢包
    }
    unsigned int end_seq = 0xFFFFFFFF;
    sendto(sock, &end_seq, 4, 0, (struct sockaddr*)&dest_addr, sizeof(dest_addr));
    (*packet_count)++;
    close(sock);
}

void send_tcp(const char *dest_ip, int dest_port, int local_port, unsigned char *buffer, size_t filesize, size_t *packet_count) {
    int sock = socket(AF_INET, SOCK_STREAM, 0);
    if (sock < 0) {
        perror("socket error");
        exit(1);
    }
    if (local_port > 0) {
        struct sockaddr_in local_addr;
        memset(&local_addr, 0, sizeof(local_addr));
        local_addr.sin_family = AF_INET;
        local_addr.sin_addr.s_addr = htonl(INADDR_ANY);
        local_addr.sin_port = htons(local_port);
        if (bind(sock, (struct sockaddr*)&local_addr, sizeof(local_addr)) < 0) {
            perror("bind local port error");
            close(sock);
            exit(1);
        }
    }
    struct sockaddr_in dest_addr;
    memset(&dest_addr, 0, sizeof(dest_addr));
    dest_addr.sin_family = AF_INET;
    dest_addr.sin_port = htons(dest_port);
    inet_pton(AF_INET, dest_ip, &dest_addr.sin_addr);
    if (connect(sock, (struct sockaddr*)&dest_addr, sizeof(dest_addr)) < 0) {
        perror("connect error");
        close(sock);
        exit(1);
    }
    // 先发送8字节文件大小
    uint64_t fsize = (uint64_t)filesize;
    send(sock, &fsize, sizeof(fsize), 0);
    size_t total_sent = 0;
    while (total_sent < filesize) {
        size_t chunk = (filesize - total_sent > CHUNK_SIZE) ? CHUNK_SIZE : (filesize - total_sent);
        ssize_t sent = send(sock, buffer + total_sent, chunk, 0);
        if (sent < 0) {
            perror("send error");
            break;
        }
        total_sent += sent;
        (*packet_count)++;
    }
    // 发送MD5校验值
    unsigned char md5[16];
    MD5_CTX ctx;
    MD5Init(&ctx);
    MD5Update(&ctx, buffer, filesize);
    MD5Final(md5, &ctx);
    send(sock, md5, 16, 0);
    close(sock);
}

int main(int argc, char *argv[]) {
    const char *file_path = "./data/test.bin";
    const char *dest_ip = DEFAULT_IP;
    int dest_port = DEFAULT_PORT;
    int local_port = 0;
    const char *protocol = "udp";
    if (argc > 1) file_path = argv[1];
    if (argc > 2) protocol = argv[2];
    if (argc > 3) dest_ip = argv[3];
    if (argc > 4) dest_port = atoi(argv[4]);
    if (argc > 5) local_port = atoi(argv[5]);

    FILE *fp = fopen(file_path, "rb");
    if (!fp) {
        perror("Failed to open file");
        return 1;
    }
    struct stat st;
    if (stat(file_path, &st) != 0) {
        perror("stat error");
        fclose(fp);
        return 1;
    }
    size_t filesize = st.st_size;
    unsigned char *buffer = (unsigned char *)malloc(filesize);
    if (!buffer) {
        perror("malloc error");
        fclose(fp);
        return 1;
    }
    fread(buffer, 1, filesize, fp);
    fclose(fp);

    struct timeval start, end;
    gettimeofday(&start, NULL);
    size_t packet_count = 0;
    if (strcmp(protocol, "udp") == 0) {
        send_udp(dest_ip, dest_port, local_port, buffer, filesize, &packet_count);
    } else if (strcmp(protocol, "tcp") == 0) {
        send_tcp(dest_ip, dest_port, local_port, buffer, filesize, &packet_count);
    } else {
        fprintf(stderr, "Unknown protocol: %s\n", protocol);
        free(buffer);
        return 1;
    }
    gettimeofday(&end, NULL);
    double elapsed = (end.tv_sec - start.tv_sec) + (end.tv_usec - start.tv_usec) / 1e6;
    double mb = filesize / (1024.0 * 1024.0);
    double speed = mb / elapsed;
    printf("Send finished. Total bytes: %zu\n", filesize);
    printf("Elapsed time: %.3f s, Speed: %.3f MB/s\n", elapsed, speed);
    printf("Total packets sent: %zu\n", packet_count);
    free(buffer);
    return 0;
}
